from typing import Optional, List, Any

from FlagEmbedding import LayerWiseFlagLLMReranker
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

import config

cpm = APIRouter(
    prefix="",
    tags=["cpm"],
    responses={404: {"description": "Not found"}},
)


class MyLayerWiseFlagLLMReranker(LayerWiseFlagLLMReranker):
    def __init__(
            self,
            model_name_or_path: str = None,
            cache_dir: str = None,
    ) -> None:
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,
                                                       cache_dir=cache_dir,
                                                       trust_remote_code=True)

        self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
                                                          cache_dir=cache_dir,
                                                          trust_remote_code=True,
                                                          local_files_only=True,
                                                          torch_dtype=torch.bfloat16,
                                                          device_map='auto')

        self.model_name_or_path = model_name_or_path
        self.cache_dir = cache_dir
        self.device = torch.device("cuda")

        self.yes_loc = self.tokenizer('Yes', add_special_tokens=False)['input_ids'][0]



# 开始启动
@cpm.on_event("startup")
def startup():
    global reranker
    reranker = MyLayerWiseFlagLLMReranker(model_name_or_path=config.cpm_reranker["model_path"])

class AloneQuery(BaseModel):
    query: str
    contexts: List[str]
    limit: Optional[float] = 0


@cpm.post("/rerank")
async def get_score(query: AloneQuery):
    result = []
    try:
        with torch.no_grad():
            pairs = [(query.query, item) for item in query.contexts]
            scores = reranker.compute_score(pairs, cutoff_layers=[28], normalize=True)
        torch.cuda.empty_cache()
        for i, score in enumerate(scores[0]):
            if score > query.limit:
                result.append({"index": i,"score": score})
        return {"code": 200, "message": "重排成功", "data": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"重排失败: {str(e)}")
    finally:
        torch.cuda.empty_cache()
